Skip to content

AIR CLI Integration: Port code_source Snapshotting#5897

Merged
riddhibhagwat-db merged 7 commits into
air-clifrom
air-integration-m4
Jul 14, 2026
Merged

AIR CLI Integration: Port code_source Snapshotting#5897
riddhibhagwat-db merged 7 commits into
air-clifrom
air-integration-m4

Conversation

@riddhibhagwat-db

@riddhibhagwat-db riddhibhagwat-db commented Jul 11, 2026

Copy link
Copy Markdown

Changes

Ports code_source snapshotting from the Python air CLI to Go and wires it into air run. When a config sets code_source.type: snapshot, the CLI now packages the local repo into a tarball, uploads it to the workspace (or a Volume), records git provenance sidecars, and attaches the paths (code_source_path, git_state_path, git_diff_path) to the ai_runtime_task payload.

Git-archive tarballs are cached by (commit, include_paths), so re-submitting the same commit skips packaging and upload.

Why

air run needs to ship the user's code to the training workload. This is a local-only port: we archive the local checkout rather than cloning from a remote, since that remote-fetch path is being deprecated. Consequences of that decision, which aren't obvious from the code alone:

  • git.remote is rejected with an actionable error (pointing users to git.commit) instead of silently fetching.
  • A ref must resolve to a commit already present locally.
  • A dirty working tree combined with git.branch is a hard error and git.branch deploys the committed HEAD, so uncommitted changes would silently be dropped.

Tests

  • Unit: cache-key parity with Python golden fixtures, git queries, the resolve decision tree, tar builders, and provenance sidecars
  • Acceptance: air run --dry-run validation & a full recorded-request submit asserting the snapshot paths land on the payload
  • End-to-end manual verification on a live workspace: exercised git-archive, plain-tar (dirty, with git_diff.patch), cache-hit reuse, and the dirty+git.branch guard; inspected the uploaded tarball and git_state.json.
Screenshot 2026-07-13 at 2 37 03 PM Screenshot 2026-07-13 at 1 56 30 PM Screenshot 2026-07-13 at 1 58 26 PM Screenshot 2026-07-13 at 3 06 59 PM Screenshot 2026-07-13 at 3 07 08 PM Screenshot 2026-07-13 at 3 07 19 PM

inspecting tar:
Screenshot 2026-07-13 at 2 42 34 PM
Screenshot 2026-07-13 at 2 51 41 PM

First slice of the code_source snapshot port (Go equivalent of the Python
CLI's cli/utils/snapshot.py). Adds the snapshot subpackage with the pure
cache-key function and a golden-fixture harness:

- cachekey.go: ComputeCacheKey + PackagingVersion="v1", ported verbatim from
  compute_snapshot_cache_key so both CLIs agree on the key algorithm.
- testdata/cache_keys.json: golden keys captured from the exact Python
  algorithm over the local-only matrix (commit + include_paths permutations).
- cachekey_test.go: byte-for-byte golden table test plus a properties test
  pinning the normalization (trim + sort, no dedup, nil == empty).

Co-authored-by: Isaac
…only)

Continues the code_source snapshot port. Flattens the snapshot files into the
aircmd package (matching the rest of experimental/air/cmd; an organizational
layer can land later) and adds Phase 2/3:

- snapshot_git.go: local, no-network git introspection ported from
  cli/utils/git_state.py — is-repo, HEAD/branch SHA, dirty check (whole-tree
  and path-scoped), commit-exists, include-path validation. Concrete gitRepo
  shelling out with arg slices (no shell string); tested against real temp
  repos, matching the Python tests.
- snapshot_resolve.go: the mode/ref decision tree (cli_entrypoint ~1541-1722),
  local-only. git.commit (must exist locally) / git.branch (local HEAD) ->
  git_archive; no ref or non-git dir -> plain_tar. Dirty check runs once and is
  threaded into the plan; dirty + git.branch is a hard error.
- snapshot_package.go: git archive / plain tar builders + .gitignore->exclude
  parsing ported from cli/utils/snapshot.py. Shells out for parity; preserves
  the top-level dir name (--prefix / -C parent) the remote extraction depends
  on, and the .git / macOS AppleDouble exclusions.
- runconfig.go: reject a truthy git.remote (remote-fetch path deprecated) with
  an actionable error pointing to git.commit; remote: false stays valid.

The remote-fetch flavor (fetch_branch_sha, partial clone, remote HEAD) is
intentionally not ported: the snapshot archives the local copy only.

Co-authored-by: Isaac
Completes the code_source snapshot port (Phases 4-5): git provenance sidecars,
the top-level orchestrator, upload via libs/filer, and the submit wiring.

- snapshot_gitstate.go: git_state.json (schema_version 1, Python field names,
  null-for-absent encoding) + git_diff.patch capture (dirty-only, 1MB cap, 5s
  timeout). Best-effort — a failure warns and continues, never fails submit.
  Merge-base/repo-url read local refs only (no fetch), matching the local-only
  design.
- snapshot.go: runSnapshot orchestrator (resolve -> package+upload ->
  sidecars). git_archive is cache-keyed and a workspace/Volume existence check
  skips packaging+upload on a hit; plain_tar is timestamp-named. Uploads through
  libs/filer (WorkspaceFilesClient, or FilesClient for a remote_volume). Also
  ports root_path resolution (project_root/, ~, env vars).
- snapshot_git.go: add currentBranch / remoteURL / mergeBaseWithUpstream and a
  runBytes variant for the raw diff bytes.
- runsubmit.go: drop the code_source rejection; package the snapshot and attach
  code_source_path / git_state_path / git_diff_path to the ai_runtime_task.

Tested end to end against the in-process fake workspace, including the cache-hit
skip and dirty-diff sidecar paths.

Co-authored-by: Isaac
Phase 6 of the snapshot port — acceptance tests over the user-visible surface.

- experimental/air/run: add dry-run cases for the new code_source config
  surface — a valid snapshot config passes validation, and a git.remote config
  is rejected with the actionable error.
- experimental/air/run-submit: a full (non-dry-run) submit that packages a git
  code_source, uploads the tarball + git_state sidecar, and asserts
  code_source_path / git_state_path land on the ai_runtime_task payload. The
  commit is date-pinned so the cache-key-derived tarball name is stable, and
  RecordRequests captures the submitted payload.

Co-authored-by: Isaac
@riddhibhagwat-db riddhibhagwat-db changed the title AIR CLI Integration: port code_source snapshotting (local-only) AIR CLI Integration: Port code_source Snapshotting Jul 11, 2026
@eng-dev-ecosystem-bot

eng-dev-ecosystem-bot commented Jul 11, 2026

Copy link
Copy Markdown
Collaborator

Integration test report

Commit: 2964af6

Run: 29297303570

Env 💚​RECOVERED 🙈​SKIP ✅​pass 🙈​skip Time
💚​ aws linux 4 4 230 1085 5:00
💚​ aws windows 4 4 232 1083 7:01
💚​ aws-ucws linux 4 4 314 1003 5:25
💚​ aws-ucws windows 4 4 316 1001 7:01
💚​ azure linux 4 4 230 1084 5:02
💚​ azure windows 4 4 232 1082 7:07
💚​ azure-ucws linux 4 4 316 1000 5:46
💚​ azure-ucws windows 4 4 318 998 8:00
💚​ gcp linux 4 4 229 1086 4:54
💚​ gcp windows 4 4 231 1084 7:28
8 interesting tests: 4 RECOVERED, 4 SKIP
Test Name aws linux aws windows aws-ucws linux aws-ucws windows azure linux azure windows azure-ucws linux azure-ucws windows gcp linux gcp windows
💚​ TestAccept 💚​R 💚​R 💚​R 💚​R 💚​R 💚​R 💚​R 💚​R 💚​R 💚​R
🙈​ TestAccept/bundle/invariant/no_drift 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S
🙈​ TestAccept/bundle/resources/vector_search_endpoints/drift/recreated_same_name 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S
🙈​ TestAccept/bundle/resources/vector_search_indexes/recreate/embedding_dimension 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S
🙈​ TestAccept/ssh/connection 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S 🙈​S
💚​ TestFetchRepositoryInfoAPI_FromRepo 💚​R 💚​R 💚​R 💚​R 💚​R 💚​R 💚​R 💚​R 💚​R 💚​R
💚​ TestFetchRepositoryInfoAPI_FromRepo/root 💚​R 💚​R 💚​R 💚​R 💚​R 💚​R 💚​R 💚​R 💚​R 💚​R
💚​ TestFetchRepositoryInfoAPI_FromRepo/subdir 💚​R 💚​R 💚​R 💚​R 💚​R 💚​R 💚​R 💚​R 💚​R 💚​R
Top 10 slowest tests (at least 2 minutes):
duration env testname
6:58 azure-ucws windows TestAccept
6:29 gcp windows TestAccept
6:10 azure windows TestAccept
6:06 aws windows TestAccept
6:00 aws-ucws windows TestAccept
2:59 azure linux TestAccept
2:54 aws-ucws linux TestAccept
2:54 gcp linux TestAccept
2:50 aws linux TestAccept
2:48 azure-ucws linux TestAccept

@ben-hansen-db

Copy link
Copy Markdown

isaac-bot review this PR

@ben-hansen-db ben-hansen-db left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

exercised git-archive, plain-tar (dirty, with git_diff.patch), cache-hit reuse, and the dirty+git.branch guard; inspected the uploaded tarball and git_state.json.

great, so commit and branch are tested?
What's the speed like for universe filtered to research folder? Can we also verify that .gitignore is being followed?

@@ -0,0 +1,2 @@
bundle:

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

something to consider is a "how to review this pr" section in your readme which can detail what are the important aspects of the pr to consider.

I'd also love to see screenshots from this working e2e

@riddhibhagwat-db riddhibhagwat-db Jul 13, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thanks for this note, i will add in a 'how to review this pr' in future PRs. I have also added screenshots from the e2e tests I ran.

Regarding the speed for the research folder snapshot:

  • git status path (resolve) takes 0.8s
  • git archive path (package & gzip) takes ~8-11s

=> total time is ~9s (warm) and ~12s (cold) , which matches the python speed

the upload path will depend on the network but it scales with the size of the tarball which is 18.9 MB (18M) (compressed from the 75MB of source)

View at: [DATABRICKS_URL]/jobs/runs/555

=== the ai_runtime_task carries the snapshot + provenance paths
>>> print_requests.py //api/2.2/jobs/runs/submit

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💯

// ai_runtime_task; a run with no code_source leaves them empty.
var snap snapshotResult
if cfg.CodeSource != nil {
snap, err = snapshotCodeSource(ctx, w, cfg.CodeSource.Snapshot, configPath, base, funcDir)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You also need to check whether cfg.CodeSource.Snapshot exists

…nce fields

The air-cli merge migrated the submit path to the typed SDK (jobs.SubmitRun /
jobs.AiRuntimeTask + w.Jobs.Submit), but left buildSubmitPayload setting
GitStatePath/GitDiffPath and referencing an undropped `snap` — neither compiles,
which is why lint, task test, and validate-generated all failed (the package
would not build).

Fix: restore the `snap` parameter, set only CodeSourcePath, and drop the two
git provenance fields from the payload. Update the code_source test to the typed
SDK types (jobs.SubmitRun / jobs.SubmitRunResponse) and stop asserting the two
dropped fields.

TEMP FIX — why this is safe, and the downstream effects:
- The typed jobs.AiRuntimeTask has no git_state_path / git_diff_path fields, and
  neither does its source proto (elastic-spark-common/.../ai_runtime_task.proto).
  So the typed SDK path cannot carry them.
- This is not a functional regression: git_state_path has ZERO consumers anywhere
  in universe, and the proto never declared the fields — so even the Python CLI's
  raw-JSON values were silently dropped on backend deserialization. Nothing ever
  read them.
- code_source_path (the field that actually delivers the code to the workload) is
  in the proto and SDK and still ships, fully typed.
- The git_state.json / git_diff.patch sidecar FILES are still uploaded next to the
  tarball (see snapshot.go) for human inspection; only the payload path pointers
  are omitted.
- Follow-up: if the backend adds git_state_path/git_diff_path to the AiRuntimeTask
  proto, regenerate the SDK and wire snap.GitStatePath/GitDiffPath back into
  buildSubmitPayload (the code spot is marked with a TODO).

Co-authored-by: Isaac
The typed-SDK submit migration changed the runs/submit request body: JSON keys
are now alphabetized and git_state_path is no longer emitted (the field was
dropped — see the buildSubmitPayload TEMP note). Regenerate the run-submit golden
output to match, and retitle the section from "snapshot + provenance paths" to
"code_source_path" since provenance paths are no longer on the payload.

Co-authored-by: Isaac
@riddhibhagwat-db riddhibhagwat-db merged commit c96d656 into air-cli Jul 14, 2026
22 checks passed
@riddhibhagwat-db riddhibhagwat-db deleted the air-integration-m4 branch July 14, 2026 01:42
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants